#!/usr/bin/python
# -*- coding: utf-8 -*-
# @Version        : 1.0
# @Update Time    : 2025/4/8 21:21
# @File           : excel_utils.py
# @IDE            : PyCharm
# @Desc           : Excel导出工具，支持大数据量导出和自定义样式

from io import BytesIO
from typing import Any, Generator, List, Dict, Optional, Union, Tuple
import xlsxwriter
from xlsxwriter.format import Format
from xlsxwriter.workbook import Workbook
from xlsxwriter.worksheet import Worksheet


class ExcelExporter:
    """Excel导出工具类，支持大数据量导出和自定义样式

    Attributes:
        title: 表格标题
        field_order: 字段顺序列表
        column_headers: 列标题列表
        summary_config: 汇总配置
        title_style: 标题样式
        header_style: 表头样式
        data_style: 数据样式
        summary_style: 汇总行样式
        max_rows_per_sheet: 每个工作表最大行数
    """

    def __init__(
        self,
        title: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        summary_config: Optional[Dict] = None,
        title_style: Optional[Dict] = None,
        header_style: Optional[Dict] = None,
        data_style: Optional[Dict] = None,
        summary_style: Optional[Dict] = None,
        max_rows_per_sheet: int = 10000,
    ):
        """初始化Excel导出器

        Args:
            title: 表格标题
            headers: 字段映射，格式为 {字段名: 列标题}
            summary_config: 汇总配置，格式为 {"columns": {字段名: "sum|avg"}}
            title_style: 标题样式配置
            header_style: 表头样式配置
            data_style: 数据样式配置
            summary_style: 汇总行样式配置
            max_rows_per_sheet: 每个工作表最大行数
        """
        self.title = title
        self.field_order = list(headers.keys()) if headers else []
        self.column_headers = list(headers.values()) if headers else []
        self.summary_config = summary_config or {}

        # 合并默认样式和自定义样式
        self.title_style = {
            "bold": True,
            "font_size": 15,
            "align": "center",
            **(title_style or {}),
        }

        self.header_style = {"bold": True, "color": "#000000", **(header_style or {})}

        self.data_style = data_style or {}
        self.summary_style = {
            "bold": True,
            "bg_color": "#C6EFCE",
            **(summary_style or {}),
        }

        self.max_rows_per_sheet = max_rows_per_sheet

    def export_to_stream(self, data: List[Dict[str, Any]]) -> BytesIO:
        """流式导出（一次性传入所有数据）

        Args:
            data: 要导出的数据列表

        Returns:
            BytesIO: 包含Excel文件的字节流
        """
        output = BytesIO()
        with xlsxwriter.Workbook(output) as workbook:
            self._process_data(workbook, data)
        output.seek(0)
        return output

    def export_from_generator(
        self, generator: Generator[Dict[str, Any], None, None]
    ) -> BytesIO:
        """流式导出（通过生成器传入数据）

        Args:
            generator: 数据生成器

        Returns:
            BytesIO: 包含Excel文件的字节流
        """
        output = BytesIO()
        with xlsxwriter.Workbook(output) as workbook:
            sheet_index = 1
            current_chunk = []

            for row in generator:
                current_chunk.append(row)
                if len(current_chunk) >= self.max_rows_per_sheet:
                    self._create_sheet(workbook, sheet_index, current_chunk)
                    sheet_index += 1
                    current_chunk = []

            if current_chunk:
                self._create_sheet(workbook, sheet_index, current_chunk)

        output.seek(0)
        return output

    def save_to_file(self, filename: str, data: List[Dict[str, Any]]) -> None:
        """保存到本地文件

        Args:
            filename: 文件名
            data: 要导出的数据列表
        """
        with xlsxwriter.Workbook(filename) as workbook:
            self._process_data(workbook, data)

    def _process_data(self, workbook: Workbook, data: List[Dict[str, Any]]) -> None:
        """处理数据并创建多个工作表

        Args:
            workbook: xlsxwriter工作簿对象
            data: 要导出的数据列表
        """
        chunks = [
            data[i : i + self.max_rows_per_sheet]
            for i in range(0, len(data), self.max_rows_per_sheet)
        ]
        for index, chunk in enumerate(chunks):
            self._create_sheet(workbook, index + 1, chunk)

    def _create_sheet(
        self,
        workbook: Workbook,
        sheet_index: int,
        data_chunk: List[Dict[str, Any]],
    ) -> None:
        """创建工作表并写入数据

        Args:
            workbook: xlsxwriter工作簿对象
            sheet_index: 工作表索引
            data_chunk: 数据块
        """
        sheet_name = f"Sheet {sheet_index}"
        sheet = workbook.add_worksheet(sheet_name)
        formats = self._create_formats(workbook)

        # 初始化行列索引
        row_idx = 0
        col_widths = [0] * len(self.column_headers) if self.column_headers else []

        # 添加标题
        if self.title:
            sheet.merge_range(
                row_idx,
                0,
                row_idx,
                len(self.column_headers) - 1,
                self.title,
                formats["title"],
            )
            row_idx += 1

        # 添加表头
        if self.column_headers:
            for col_idx, header in enumerate(self.column_headers):
                sheet.write(row_idx, col_idx, header, formats["header"])
                col_widths[col_idx] = max(col_widths[col_idx], len(str(header)))
            row_idx += 1

        # 写入数据
        for data_row in data_chunk:
            for col_idx, field_name in enumerate(self.field_order):
                cell_value = data_row.get(field_name, "")
                cell_val, fmt = self._process_cell_value(
                    workbook, cell_value, formats["data"]
                )
                sheet.write(row_idx, col_idx, cell_val, fmt)
                col_widths[col_idx] = max(
                    col_widths[col_idx],
                    len(str(cell_val)) * 1.2,
                )
            row_idx += 1

        # 添加汇总行
        if self.summary_config and data_chunk:
            self._add_summary_row(sheet, row_idx, data_chunk, formats["summary"])
            row_idx += 1

        # 设置自适应列宽
        for col_idx, width in enumerate(col_widths):
            sheet.set_column(col_idx, col_idx, width)

    def _process_cell_value(
        self, workbook: Workbook, cell_value: Any, default_format: Format
    ) -> Tuple[Any, Format]:
        """处理单元格值和样式

        Args:
            workbook: xlsxwriter工作簿对象
            cell_value: 单元格值
            default_format: 默认格式

        Returns:
            Tuple[Any, Format]: 处理后的值和格式
        """
        if isinstance(cell_value, tuple) and len(cell_value) == 2:
            cell_val, cell_style = cell_value
            fmt = workbook.add_format(cell_style)
        else:
            cell_val = cell_value
            fmt = default_format
        return cell_val, fmt

    def _add_summary_row(
        self,
        sheet: Worksheet,
        row_idx: int,
        data: List[Dict[str, Any]],
        fmt: Format,
    ) -> None:
        """添加汇总行

        Args:
            sheet: xlsxwriter工作表对象
            row_idx: 行索引
            data: 数据列表
            fmt: 汇总行格式
        """
        summary_values = []
        for field_name in self.field_order:
            if field_name in self.summary_config.get("columns", {}):
                op = self.summary_config["columns"][field_name]
                values = []
                for row in data:
                    try:
                        val = float(row.get(field_name, 0))
                        values.append(val)
                    except (ValueError, TypeError):
                        continue

                if op.lower() == "sum" and values:
                    summary_val = sum(values)
                elif op.lower() == "avg" and values:
                    summary_val = sum(values) / len(values)
                else:
                    summary_val = ""
                summary_values.append(summary_val)
            else:
                summary_values.append("")

        for col_idx, value in enumerate(summary_values):
            sheet.write(row_idx, col_idx, value, fmt)

    def _create_formats(self, workbook: Workbook) -> Dict[str, Format]:
        """创建工作簿格式

        Args:
            workbook: xlsxwriter工作簿对象

        Returns:
            Dict[str, Format]: 格式字典
        """
        return {
            "title": workbook.add_format(self.title_style),
            "header": workbook.add_format(self.header_style),
            "data": workbook.add_format(self.data_style),
            "summary": workbook.add_format(self.summary_style),
        }


# from fastapi import FastAPI
# from fastapi.responses import StreamingResponse
#
# app = FastAPI()
#
#
# @app.get("/export")
# async def export_data():
#     # 你的数据准备逻辑
#     headers = {"id": "Employee ID", "name": "Full Name", "sales": "Total Sales"}
#
#     def data_generator():
#         for i in range(1, 1000):
#             yield {"id": i, "name": f"Employee {i}", "sales": i * 100}
#
#
#     # 初始化导出器
#     exporter = ExcelExporter(
#         title="导出测试",
#         headers=headers,
#         summary_config={"columns": {"sales": "sum"}},
#         header_style={"bg_color": "#FFA07A", "align": "center"},
#         data_style={"align": "center"},
#     )
#
#     # 导出到文件
#     return StreamingResponse(
#         exporter.export_from_generator(data_generator()),
#         media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
#         headers={"Content-Disposition": "attachment; filename=report.xlsx"},
#     )
#
#
# # 使用示例
# if __name__ == "__main__":
#     import uvicorn
#
#     uvicorn.run(app, host="0.0.0.0", port=8000)
